/*-------------------------------------------
                Includes
-------------------------------------------*/
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>
#include <dlfcn.h>
#include <vector>
#include <string>
#include <stdbool.h>
#include <timer.h>
#include <queue>

#include "yolo_inference.h"
#include "log.h"
/*-------------------------------------------
                defines
-------------------------------------------*/
#ifdef LOG_TAG
#undef LOG_TAG
#endif
#define LOG_TAG "yolo_inference.cc"

#define MAX_DETECT_RESULT_GROUP_QUEUE_SIZE 10
/*-------------------------------------------
                variables
-------------------------------------------*/
#define INDENT "    "
MODEL_INFO m_info;
YOLO_INFO y_info;
LETTER_BOX letter_box;
// 用于存储检测结果的列队
std::queue<detect_result_group_t> detect_result_group_queue;
/*-------------------------------------------
                  Functions
-------------------------------------------*/

void query_dfl_len(MODEL_INFO *m, YOLO_INFO *y_info)
{
    // set dfl_len
    if ((y_info->m_type == YOLOV8) || (y_info->m_type == PPYOLOE_PLUS) || (y_info->m_type == YOLOV6))
    {
        if (m->n_output > 6)
        {
            y_info->score_sum_available = true;
        }
        else
        {
            y_info->score_sum_available = false;
        }

        if ((y_info->m_type == YOLOV8) || (y_info->m_type == PPYOLOE_PLUS))
        {
            y_info->dfl_len = (int)(m->out_attr[0].dims[1] / 4);
        }

        if (y_info->m_type == YOLOV6)
        {
            // dump_tensor_attr(&m->out_attr[0]);
            if (m->out_attr[0].dims[1] != 4)
            {
                y_info->dfl_len = (int)(m->out_attr[0].dims[1] / 4);
            }
        }
    }
}

char LABEL_TXT_PATH[64] = {0};
void yolo_init(char *model_path, char *yolo_type, const char *label_txt_path, int class_num, float conf_thresh, float nms_thresh)
{
    const char *anchor_path;
    TIMER timer;
    timer.indent_set("    ");
    int ret;

    LOG_INFO("Yolo init...\n");
    // 设置模型参数
    init_yolo_param(class_num, nms_thresh, conf_thresh);
    // 读取标签文件
    strcpy(LABEL_TXT_PATH, label_txt_path);
    // yolo模型类型
    y_info.m_type = string_to_model_type(yolo_type);
    // 指定锚点文件
    switch (y_info.m_type)
    {
    case YOLOV5:
        y_info.anchor_per_branch = 3;
        anchor_path = "./models/anchors_yolov5.txt";
        break;
    case YOLOV7:
        y_info.anchor_per_branch = 3;
        anchor_path = "./models/anchors_yolov7.txt";
        break;
    case YOLOX:
        y_info.anchor_per_branch = 1;
        break;
    default:
        y_info.anchor_per_branch = 1;
        break;
    }
    // 读取锚点
    if ((y_info.m_type == YOLOV5) || (y_info.m_type == YOLOV7))
    {
        int n = 2 * 3 * y_info.anchor_per_branch;
        LOG_INFO("%sAnchors: ", INDENT);
        float result[n];
        int valid_number;
        ret = readFloats(anchor_path, &result[0], n, &valid_number);
        for (int i = 0; i < valid_number; i++)
        {
            y_info.anchors[i] = (int)result[i];
            printf("%d ", y_info.anchors[i]);
        }
        printf("\n");
    }
    else
    {
        LOG_INFO("Anchor free\n");
    }
    // 模型post输出
    y_info.post_type = Q8;
    // 单张输入
    y_info.in_source = SINGLE_IMG;

    // 传入rknn模型路径
    m_info.m_path = model_path;
    rkdemo_init(&m_info);
    rkdemo_init_input_buffer_all(&m_info, ZERO_COPY_API, RKNN_TENSOR_UINT8);
    rkdemo_init_output_buffer_all(&m_info, ZERO_COPY_API, 0);
    query_dfl_len(&m_info, &y_info);
    // letter_box缩放处理参数
    letter_box.target_height = m_info.in_attr[0].dims[1];
    letter_box.target_width = m_info.in_attr[0].dims[2];

    // 加载数据到rknn内存
    timer.tik();
    for (int i = 0; i < m_info.n_input; i++)
    {
        rknn_set_io_mem(m_info.ctx, m_info.input_mem[i], &m_info.in_attr_native[i]);
    }
    for (int i = 0; i < m_info.n_output; i++)
    {
        rknn_set_io_mem(m_info.ctx, m_info.output_mem[i], &m_info.out_attr_native[i]);
    }
    timer.tok();
    LOG_INFO("rknn_set_io_mem(include input/output) use time: %f ms\n", timer.get_time());
}

detect_result_group_t yolo_infer_frame(unsigned char *inputData, int input_height, int input_width)
{
    void *output_buf_list[m_info.n_output];
    TIMER timer;
    timer.indent_set("    ");
    int ret;

    // Letter box resize
    if ((input_height == letter_box.target_height) && (input_width == letter_box.target_width))
    {
        // m_info.inputs[0].buf = input_data;
        LOG_DEBUG("No need to letter box resize, directly copy\n");
        memcpy(m_info.input_mem[0]->virt_addr, inputData, input_height * input_width * 3);
    }
    else
    {
        timer.tik();
        letter_box.in_height = input_height;
        letter_box.in_width = input_width;
        compute_letter_box(&letter_box);
        ret = rga_letter_box_resize(inputData, m_info.input_mem[0]->fd, &letter_box);

        if (ret != 0)
        {
            LOG_WARN("RGA letter box resize failed, use stb to resize\n");
            stb_letter_box_resize(inputData, (unsigned char *)m_info.input_mem[0]->virt_addr, letter_box);
        }
        letter_box.reverse_available = true;
        timer.tok();
        // timer.print_time("letter box");
        LOG_DEBUG("RGA letter box resize use time: %f ms\n", timer.get_time());
    }

    LOG_DEBUG("RUN MODEL\n");
    // rknn run
    timer.tik();
    ret = rknn_run(m_info.ctx, NULL);
    timer.tok();
    // timer.print_time("rknn_run");
    LOG_DEBUG("rknn_run use time: %f ms\n", timer.get_time());

    /* Post process */
    detect_result_group_t detect_result_group;
    timer.tik();
    for (int i = 0; i < m_info.n_output; i++)
    {
        // output_buf_list[i] = m_info.outputs[i].buf;
        output_buf_list[i] = m_info.output_mem[i]->virt_addr;
    }
    post_process(output_buf_list, &m_info, &y_info, &detect_result_group, (const char*)LABEL_TXT_PATH);
    timer.tok();
    // timer.print_time("cpu_post_process");
    LOG_DEBUG("cpu_post_process use time: %f ms\n", timer.get_time());

    char score_result[64];
    for (int i = 0; i < detect_result_group.count; i++)
    {
        detect_result_group.results[i].box.left = w_reverse(detect_result_group.results[i].box.left, letter_box);
        detect_result_group.results[i].box.right = w_reverse(detect_result_group.results[i].box.right, letter_box);
        detect_result_group.results[i].box.top = h_reverse(detect_result_group.results[i].box.top, letter_box);
        detect_result_group.results[i].box.bottom = h_reverse(detect_result_group.results[i].box.bottom, letter_box);

        detect_result_t *det_result = &(detect_result_group.results[i]);
        LOG_DEBUG("%s%s @ (%d %d %d %d) %f\n", INDENT,
               det_result->name,
               det_result->box.left, det_result->box.top, det_result->box.right, det_result->box.bottom,
               det_result->prop);
    }

    ret = rknn_outputs_release(m_info.ctx, m_info.n_output, m_info.outputs);
    return detect_result_group;
}


void yolo_deinit()
{
    rkdemo_release(&m_info);
}


// 将一个yolo_infer_frame的结果压入列队
void push_detect_result_group(detect_result_group_t detect_result_group)
{
    // 将detect_result_group压入列队
    detect_result_group_queue.push(detect_result_group);
    // 如果列队长度超过了最大值，就弹出最早的一条
    if (detect_result_group_queue.size() > MAX_DETECT_RESULT_GROUP_QUEUE_SIZE)
    {
        detect_result_group_queue.pop();
    }
}

// 从列队中取出一条detect_result_group
detect_result_group_t pop_detect_result_group()
{
    detect_result_group_t detect_result_group;
    // 如果列队不为空，就取出一条
    if (!detect_result_group_queue.empty())
    {
        detect_result_group = detect_result_group_queue.front();
        detect_result_group_queue.pop();
    }
    // 如果列队为空，就返回空的detect_result_group
    else
    {
        detect_result_group.count = 0;
    }
    return detect_result_group;
}




